import numpy as np
import torch
from scipy.optimize import linprog

from algorithms.algorithm import Algorithm
from individual.population import Population
from individual.solution import Solution
from problems.problem import Problem


class LZUI(Algorithm):
    def __init__(self, problem: Problem,
                 pop_size: int,
                 kwargs,
                 ):
        super().__init__(pop_size, problem, kwargs)
        self.population = self.problem.init_population(pop_size)

    def each_iteration(self):
        for i in range(self.pop_size):
            w = self.cal_w(self.population, self.population[i])
            pop_obj_g = self.calg(self.population[i])
            offspring_dec = torch.tensor(self.calf(w, pop_obj_g, self.problem.low_limit, self.problem.high_limit),
                                         dtype=torch.double)
            offspring_dec = self.problem.repair_decision(offspring_dec)
            offspring_obj = self.problem.estimate_population(offspring_dec).obj
            self.population[i] = Solution.solution_build(offspring_dec, offspring_obj)

    def calg(self, solution):
        return self.problem.cal_grad(solution.dec)

    @staticmethod
    def cal_w(population: Population, x):
        f = x.obj.numpy()
        A = population.obj.numpy()
        b = np.ones(A.shape[0])
        res = linprog(f, A_ub=-A, b_ub=-b, bounds=(0, None), method='highs')
        return res.x

    @staticmethod
    def calf(w, pop_obj_g, lower, upper):
        f = (pop_obj_g * w).sum(axis=1)
        res = linprog(f, bounds=np.vstack((lower, upper)).T)
        return res.x
